# viz/tree.py —— 表血缘（目标 <- 来源）树状 with subquery grouping & coloring
from typing import List, Dict, Any, Iterable, Optional, Set
from collections import defaultdict

from pyecharts import options as opts
from pyecharts.charts import Tree
from core.models import StatementLineage


def _node_color(name: str) -> Optional[str]:
    n = name.lower()
    if n.startswith("hive."):
        return "#8be9fd"   # 青
    if n.startswith("mysql"):
        return "#bd93f9"   # 紫
    if n.startswith("subquery"):
        return "#50fa7b"   # 绿
    return None


def _mk_node(name: str, children: Optional[List[Dict[str, Any]]] = None) -> Dict[str, Any]:
    node = {"name": name}
    color = _node_color(name)
    if color:
        node["itemStyle"] = {"color": color}
    if children:
        node["children"] = children
    return node


def _group_by_schema(tables: Iterable[str]) -> List[Dict[str, Any]]:
    """把 catalog.schema.table 分组为 schema 节点，便于折叠。"""
    bucket: Dict[str, List[str]] = defaultdict(list)
    for t in tables:
        parts = t.split(".")
        if len(parts) >= 2:
            key = ".".join(parts[:-1])  # catalog.schema
        else:
            key = "(default)"
        bucket[key].append(t)

    children: List[Dict[str, Any]] = []
    for sch, tbls in sorted(bucket.items()):
        children.append(
            _mk_node(sch, children=[_mk_node(tt) for tt in sorted(tbls)])
        )
    return children


def render_table_tree(
    st: StatementLineage,
    *,
    title: str = "表血缘",
    group_by_schema: bool = True,
    width: str = "1000px",
    height: str = "680px",
    return_html: bool = True,
):
    """
    - SELECT: 根 = "SELECT"
        - 子节点1：每个 Subquery <alias>，其 children = 该子查询内的来源表（可分组）
        - 子节点2：直接来源表（不属于任何子查询）
    - 写入类：每个目标表为根，子节点同上（来源表/子查询）
    - 节点按来源系统着色（Hive/MySQL/Subquery）
    """
    # 子查询映射（extractor 已写入 ctes['_subq']）
    subq_map: Dict[str, List[str]] = st.ctes.get("_subq", {}) if isinstance(st.ctes, dict) else {}
    subq_aliases = sorted(subq_map.keys())
    subq_tables_union: Set[str] = set()
    for v in subq_map.values():
        subq_tables_union |= set(v)

    direct_sources = sorted(set(st.source_tables) - subq_tables_union)

    def _children_for_sources(tables: Iterable[str]) -> List[Dict[str, Any]]:
        if group_by_schema:
            return _group_by_schema(tables)
        return [_mk_node(t) for t in sorted(tables)]

    def _subquery_nodes() -> List[Dict[str, Any]]:
        nodes: List[Dict[str, Any]] = []
        for alias in subq_aliases:
            tables = subq_map.get(alias, [])
            label = f"Subquery {alias}"
            nodes.append(_mk_node(label, _children_for_sources(tables)))
        return nodes

    is_select = (st.op.upper() == "SELECT")
    roots = list(st.target_tables)

    data: List[Dict[str, Any]] = []
    if is_select or not roots:
        root_name = "SELECT"
        kids = []
        kids.extend(_subquery_nodes())
        if direct_sources:
            kids.extend(_children_for_sources(direct_sources))
        data = [_mk_node(root_name, kids)]
    else:
        # 非 SELECT：每个目标表作为一棵树的根
        for tgt in sorted(roots):
            kids = []
            kids.extend(_subquery_nodes())
            if direct_sources:
                kids.extend(_children_for_sources(direct_sources))
            data.append(_mk_node(tgt, kids))

    chart = (
        Tree(init_opts=opts.InitOpts(width=width, height=height))
        .add(
            series_name="",
            data=data,
            orient="LR",
            symbol="emptyCircle",
            label_opts=opts.LabelOpts(position="left"),
        )
        .set_global_opts(title_opts=opts.TitleOpts(title=f"{title} - {st.op}"))
    )
    return chart.render_embed() if return_html else chart
